# This code is part of cqlib.
#
# Copyright (C) 2025 China Telecom Quantum Group, QuantumCTek Co., Ltd.,
# Center for Excellence in Quantum Information and Quantum Physics.
#
# This code is licensed under the Apache License, Version 2.0. You may
# obtain a copy of this license in the LICENSE file in the root directory
# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
#
# Any modifications or derivative works of this code must retain this
# copyright notice, and modified files need to carry a notice indicating
# that they have been altered from the originals.

"""
Text-based quantum circuit visualization module

Provides ASCII/Unicode art representation of quantum circuits using box-drawing characters.
"""
import copy
import logging
import shutil

from cqlib.circuits.circuit import Circuit
from cqlib.circuits.gates import SWAP, CZ
from cqlib.circuits.barrier import Barrier
from cqlib.circuits.gates.gate import ControlledGate
from cqlib.circuits.instruction_data import InstructionData
from cqlib.circuits.qubit import Qubit

from .base import BaseDrawer, BoxChar

logger = logging.getLogger('cqlib.vis')


class TextDrawer(BaseDrawer):
    """Render quantum circuits as text diagrams using box-drawing characters."""

    def __init__(
            self,
            circuit: Circuit,
            qubit_order: list[int | Qubit] | None = None,
            line_width: int | None = None,
            initial_state: bool = False,
            reverse_bits: bool = False,
    ):
        """Initialize a text drawer instance.

        Args:
            circuit (Circuit): Quantum circuit to visualize.
            qubit_order (list[int | Qubit] | None): Custom ordering of qubits.
            line_width (int | None): Optional width of the visualization output.
            initial_state (bool): Whether to display initial state labels for qubits.
            reverse_bits (bool): Whether to reverse the qubit order in display.
        """
        super().__init__(circuit, qubit_order)
        self.line_width = line_width
        self.initial_state = initial_state
        self.reverse_bits = reverse_bits
        sorted_qubits = self.sorted_qubits
        if reverse_bits and qubit_order is None:
            sorted_qubits = list(reversed(sorted_qubits))
        
        self.display_qubits = sorted_qubits

    def __str__(self):
        """Return the rendered circuit as a user-friendly string."""
        return self.drawer()

    def __repr__(self):
        """Return the rendered circuit for debugging contexts."""
        return self.drawer()

    def drawer(self) -> str:
        """Generate the complete text-based circuit diagram.

        Returns:
            str: Full circuit representation assembled from box-drawing characters.
        """
        max_line_width = self.get_line_width()
        lines = self.make_lines()
        data = []
        lines_count = len(self.display_qubits) * 2 + 1
        start_qubits = copy.deepcopy(lines[0])
        current_data = [lines[0][i] for i in range(lines_count)]
        current_width = 0
        for moment in lines[1:]:

            moment_len = len(''.join(moment[0]))
            if moment_len + current_width > max_line_width:
                data.append(current_data)
                for d in current_data:
                    d.append(BoxChar.RIGHT_ARROW)
                current_data = []
                for i in range(lines_count):
                    s = [BoxChar.LEFT_ARROW]
                    s.extend(start_qubits[i])
                    current_data.append(s)
                current_width = 0
            for i, line in enumerate(moment):
                current_data[i].extend(line)
            current_width += moment_len
        data.append(current_data)
        t = []
        for lines in data:
            for line in lines:
                t.extend(line)
                t.append('\n')
            t.append('\n')
        return ''.join(t)

    # pylint: disable=too-many-locals
    def make_lines(self) -> list[list[list[str]]]:
        """Construct the circuit visualization line structure.

        Returns:
            list[list[list[str]]]: Nested list representation of the circuit.
                The outer list contains header/moment segments, each middle list
                captures per-line content (labels and wires), and inner lists hold
                character chunks.
        """
        lines = []
        qubit_len = max(len(str(q.index)) for q in self.display_qubits)
        lines_count = len(self.display_qubits) * 2 + 1

        # initial state
        if self.initial_state:
            initial_value = "|0>"
            empty_line = ' ' * (qubit_len + 9)
        else:
            initial_value = ""
            empty_line = ' ' * (qubit_len + 6)

        # qubits
        start_lines = []
        for qubit in self.display_qubits:
            start_lines.extend([
                [empty_line],
                [''.join([' ', f'Q{qubit.index}'.rjust(qubit_len + 1),
                          ': ', initial_value, BoxChar.LEFT_RIGHT * 2])]
            ])
        start_lines.append([empty_line])
        lines.append(start_lines)

        # instructions
        for moment in self.generate_moment():
            # before moment, add one symbol
            moment_lines = [[BoxChar.LEFT_RIGHT if i % 2 == 1 else ' '] for i in range(lines_count)]

            # drawer one moment, maybe multi columns
            columns = self.moment_to_columns(moment)
            for column in columns:
                column_lines = self.draw_column(column)
                for line_i, line in enumerate(column_lines):
                    moment_lines[line_i].extend(line)

            # after moment, add one symbol
            for i, line in enumerate(moment_lines):
                line.append(BoxChar.LEFT_RIGHT if i % 2 == 1 else ' ')

            # mark many column as one moment
            col_len = len(columns)
            if col_len > 1:
                s = BoxChar.LEFT_RIGHT * (len(''.join(moment_lines[0])) - 2)
                moment_lines[0] = [BoxChar.BOTTOM_RIGHT, s, BoxChar.BOTTOM_LEFT]
                moment_lines[-1] = [BoxChar.TOP_RIGHT, s, BoxChar.TOP_LEFT]
            lines.append(moment_lines)
        return lines

    def draw_column(self, column: list[InstructionData]):
        """Process a vertical column of parallel operations.

        Args:
            column (list[InstructionData]): Non-overlapping operations from the same moment.

        Returns:
            list[list[str]]: Formatted line buffers ready for insertion into the drawing.
        """
        max_width = 1

        # container
        lines = [[] for _ in range(len(self.display_qubits) * 2 + 1)]
        # draw every InstructionData
        for ins in column:
            match len(ins.qubits):
                case 1:
                    lines = self.draw_single_gate(ins, lines)
                case _:
                    lines = self.draw_multi_gate(ins, lines)
        # calculate max width
        for line in lines:
            for s in line:
                max_width = max(max_width, len(s))
        # fit max width
        empty_line = ' ' * max_width
        left_right_line = BoxChar.LEFT_RIGHT * max_width
        for i, line in enumerate(lines):
            if i % 2 == 0:
                if line:
                    lines[i] = [line[0].center(max_width)]
                else:
                    lines[i].append(empty_line)
            else:
                if line:
                    lines[i] = [line[0].center(max_width, BoxChar.LEFT_RIGHT)]
                else:
                    lines[i].append(left_right_line)

        return lines

    def draw_single_gate(self, ins: InstructionData, lines: list[list[str]]):
        """Render a single-qubit gate into the line buffers.

        Args:
            ins (InstructionData): Gate instruction to render.
            lines (list[list[str]]): Current drawing line buffers.

        Returns:
            list[list[str]]: Updated line buffers with the gate symbol inserted.
        """

        if isinstance(ins.instruction, Barrier):
            idx = self.qubit_line_no(ins.qubits[0])
            lines[idx].append(BoxChar.TOP_BOTTOM)
            return lines

        lines[self.qubit_line_no(ins.qubits[0])].append(self._str_instruction(ins.instruction))
        return lines

    def draw_multi_gate(self, ins: InstructionData, lines: list[list[str]]):
        """Render a multi-qubit operation with vertical connectors.

        Handles SWAP, CZ, barrier, and general controlled gates using appropriate
        box-drawing characters.

        Args:
            ins (InstructionData): Multi-qubit instruction to render.
            lines (list[list[str]]): Current drawing line buffers.

        Returns:
            list[list[str]]: Updated line buffers with gate symbols and connectors.
        """
        if isinstance(ins.instruction, SWAP):
            lines[self.qubit_line_no(ins.qubits[0])].append(BoxChar.CONNECT)
            lines[self.qubit_line_no(ins.qubits[1])].append(BoxChar.CONNECT)
        elif isinstance(ins.instruction, CZ):
            lines[self.qubit_line_no(ins.qubits[0])].append(BoxChar.DOT)
            lines[self.qubit_line_no(ins.qubits[1])].append(BoxChar.DOT)
        elif isinstance(ins.instruction, Barrier):
            for qubit in ins.qubits:
                idx = self.qubit_line_no(qubit)
                if idx not in (1, len(lines) - 1):
                    lines[idx - 1].append(BoxChar.TOP_BOTTOM)
                lines[idx].append(BoxChar.TOP_BOTTOM)
            # Barrier, No vertical connections required
            return lines
        elif isinstance(ins.instruction, ControlledGate):
            for index, qubit in enumerate(ins.qubits):
                if index in ins.instruction.control_index:
                    lines[self.qubit_line_no(qubit)].append(BoxChar.DOT)
                else: 
                    lines[self.qubit_line_no(qubit)].append(self._str_instruction(ins.instruction.base_gate))

        # add connect vertical line
        qubit_index = [self.qubit_line_no(q) for q in ins.qubits]
        min_index, max_index = min(qubit_index), max(qubit_index)
        for idx in range(min_index + 1, max_index):
            lines[idx].append(BoxChar.TOP_BOTTOM_LEFT_RIGHT if idx % 2 == 1
                              else BoxChar.TOP_BOTTOM)

        return lines

    def get_line_width(self):
        """Determine the effective display width used for rendering.

        Prefers an explicit ``line_width``, then the current terminal width, and
        finally a fixed default.

        Returns:
            int: Display width in characters.
        """
        if self.line_width:
            if self.line_width > self.MIN_LINE_WIDTH:
                return self.line_width
            if self.line_width < 0:
                return float('inf')
        # try to get terminal size
        width, _ = shutil.get_terminal_size()
        if not width:
            width = self.DEFAULT_LINE_WIDTH
        return width
    
    def qubit_line_no(self, qubit: Qubit) -> int:
        """Return the line index for a qubit in the current display order.

        Overrides the base implementation to use ``display_qubits`` rather than
        ``sorted_qubits``.

        Args:
            qubit (Qubit): Qubit whose row index is requested.

        Returns:
            int: Odd-numbered line index where the qubit is rendered.

        Raises:
            ValueError: If the qubit is not included in the display order.
        """
        try:
            index = self.display_qubits.index(qubit)
            return index * 2 + 1
        except ValueError as exc:
            raise ValueError(f"Qubit {qubit} not found in display order") from exc
    
    
    
def draw_text(
        circuit: Circuit,
        qubit_order: list[int | Qubit] = None,
        line_width: int = None,
        initial_state: bool = False,
        reverse_bits: bool = False,
):
    """Render a circuit using the text visualization backend.

    Args:
        circuit (Circuit): Quantum circuit to visualize.
        qubit_order (list[int | Qubit] | None): Optional custom qubit arrangement.
        line_width (int | None): Optional output width; ``None`` auto-detects terminal width.
        initial_state (bool): Whether to include initial state labels.
        reverse_bits (bool): Whether to reverse the display order.

    Returns:
        str: Ready-to-print circuit diagram.
    """
    return TextDrawer(circuit, qubit_order, line_width, initial_state, reverse_bits)
